# Copyright (c) Open-MMLab. All rights reserved.
from torch.nn.utils import clip_grad

from .hook import Hook

import torch
from apex import amp

class OptimizerHook(Hook):

    def __init__(self, grad_clip=None):
        self.grad_clip = grad_clip

    def clip_grads(self, params):
        clip_grad.clip_grad_norm_(
            filter(lambda p: p.requires_grad, params), **self.grad_clip)

    def after_train_iter(self, runner):

        # with torch.autograd.profiler.profile(record_shapes=True, use_npu=True) as prof:
        runner.optimizer.zero_grad()

        # fox:设置混合精度
        with amp.scale_loss(runner.outputs['loss'], runner.optimizer) as scaled_loss:
            scaled_loss.backward()

            # prof.table(row_limit=200000)

        if self.grad_clip is not None:
            self.clip_grads(runner.model.parameters())
        runner.optimizer.step()
